﻿// ----------
// AZPEP SysRi Tool
// Copyright (c) AZPEP_BiliBili(@零分的考卷).
// All Rights Reserved.
//
// ----------
// 文件名：CreatePE.cs
//


//using Make_WinPE;
//using Sunny.UI.Win32;
using AZPEP_SysRi_Tool;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Reflection.Metadata;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

#pragma warning disable CA1416

namespace AZPEP_SysRi_Tool
{
    internal class CreatePE
    {
        // 创建PE启动项所用的两组GUID，请勿乱改！
        private const string GUID1 = "{ffffffff-8d96-11de-8e71-ffffffffffff}";
        private const string GUID2 = "{ffffffff-8d96-11de-8e71-fffffffffffe}";

        /// <summary>
        /// BcdBoot文件路径
        /// </summary>
        private static readonly string BcdEditFilePath = $"{System.Environment.GetEnvironmentVariable("SystemRoot")}\\System32\\bcdedit.exe";

        /// <summary>
        /// 制作PE启动盘
        /// </summary>
        /// <param name="wimFilePath">PE wim 文件绝对路径</param>
        /// <param name="sdiFilePath">PE SDI 文件绝对路径</param>
        /// 
        /// <param name="ImageFilePath">映像文件路径</param>
        /// <param name="ImageIndex">欲释放映像索引号</param>
        /// 
        /// <param name="ReInstallPartitionLetter">重装系统分区盘符，应如“C”而不是“C:\”</param>
        /// <param name="EfiPartitionLetter">ESP分区盘符，应如“C”而不是“C:\”</param>
        /// 
        /// <param name="DisableWindowsDefender">禁用“Windows Defender”</param>
        /// <param name="DisableWindowsUpdate">禁用Windows自动更新</param>
        /// <param name="SkipMicrosoftLogin">跳过Microsoft登陆验证</param>
        /// 
        /// <param name="ReturnCode">返回错误代码，成功则返回SUCCESS</param>
        /// <returns></returns>
        public static bool MakePEStartupEntry(
            string wimFilePath,
            string sdiFilePath,

            string ImageFilePath,
            int ImageIndex,

            string EfiPartitionLetter,
            string ReInstallPartitionLetter,

            long SystemSizeMB,

            bool DisableWindowsDefender,
            bool DisableWindowsUpdate,
            bool SkipMicrosoftLogin,

            out ErrCode.ErrorCode ReturnCode
            )
        {
            LogProtocol.WriteLog.Info("[CREATE_PE]   正在检查并删除过时的必要文件...");

            /* 删除存在的必要文件 */
            DriveInfo[] DIL = DriveInfo.GetDrives();
            foreach (DriveInfo DI in DIL)
            {
                if (!DI.IsReady) continue;
                if (File.Exists($"{DI.Name}AZPEP_SysRi_Tool__ReInstallPartition"))
                {
                    LogProtocol.WriteLog.Warning($"[CREATE_PE]   文件\"{DI.Name}AZPEP_SysRi_Tool__ReInstallPartition\"已存在，正在删除...");
                    File.Delete($"{DI.Name}AZPEP_SysRi_Tool__ReInstallPartition");
                }
                //LogProtocol.WriteLog.Info($"[CREATE_PE][DELETE]   文件\"{DI.Name}AZPEP_SysRi_Tool__ReInstallPartition\"不存在或已删除");

                if (File.Exists($"{DI.Name}AZPEP_SysRi__MainProgramPartition"))
                {
                    LogProtocol.WriteLog.Warning($"[CREATE_PE]   文件\"{DI.Name}AZPEP_SysRi__MainProgramPartition\"已存在，正在删除...");
                    File.Delete($"{DI.Name}AZPEP_SysRi__MainProgramPartition");
                }
                //LogProtocol.WriteLog.Info($"[CREATE_PE][DELETE]   文件\"{DI.Name}AZPEP_SysRi_Tool__MainProgramPartition\"不存在或已删除");

                if (File.Exists($"{DI.Name}AZPEP_SysRi_Tool__EfiPartition"))
                {
                    LogProtocol.WriteLog.Warning($"[CREATE_PE]   文件\"{DI.Name}AZPEP_SysRi__EfiPartition\"已存在，正在删除...");
                    File.Delete($"{DI.Name}AZPEP_SysRi_Tool__EfiPartition");
                }
                //LogProtocol.WriteLog.Info($"[CREATE_PE][DELETE]    文件\"{DI.Name}AZPEP_SysRi_Tool_EfiPartition\"不存在或已删除");

                if (File.Exists($"{DI.Name}AZPEP_SysRi_Tool__ShrinkPartition"))
                {
                    LogProtocol.WriteLog.Warning($"[CREATE_PE]   文件\"{DI.Name}AZPEP_SysRi__ShrinkPartition\"已存在，正在删除...");
                    File.Delete($"{DI.Name}AZPEP_SysRi_Tool__ShrinkPartition");
                }
                //LogProtocol.WriteLog.Info($"[CREATE_PE][DELETE]    文件\"{DI.Name}AZPEP_SysRi__ShrinkPartition\"不存在或已删除");
            }

            /* 检查PE文件是否能够正常打开 */
            /* 这部我已经不想做了，一是时间不够，二是我相信用户基本没有这个实力 */
            /* 所以就我们跳过这一步吧 */


            /* 计算分区大小 */
            /* 我不确定使用long而不是decimal是否会产生精度误差 */

            LogProtocol.WriteLog.Info("[CREATE_PE]   正在计算分区大小...");
            long wimFileSize = (long)Math.Round(new FileInfo(wimFilePath).Length * 1.2);
            long sdiFileSize = new FileInfo(sdiFilePath).Length + 16384;
            long imageFileSize = (long)Math.Round(new FileInfo(ImageFilePath).Length * 1.15);

            long MinPartitionSize = wimFileSize + sdiFileSize + imageFileSize + (16 * 1024 ^ 2); //16*1024^2 会给分区留下至少16MB的保留空间

            LogProtocol.WriteLog.Info($"[CREATE_PE]   分区最小大小：{MinPartitionSize} 字节");

            /* 创建PE分区 */
            char? MainProgramPartitionLetter = null; //分区盘符
            char? ShrinkPartitionLetter = null; //被压缩卷盘符，默认没有

            ErrCode.ErrorCode CallBackErr = ErrCode.ErrorCode.SUCCESS;

            bool Status = DiskPart.CreatePEDisk(
                ReInstallPartitionLetter[0],
                SystemSizeMB,

                MinPartitionSize / 1024 / 1024,

                out MainProgramPartitionLetter,
                out ShrinkPartitionLetter,

                out CallBackErr
                );

            if (!Status)
            {
                //创建失败，日志报错，输出错误代码，返回

                //处理错误信息

                string ErrInfo = "";
                string ErrCode2 = "";
                switch (CallBackErr)
                {
                    case ErrCode.ErrorCode.__MAINPROG__CREATEPE_NO_MORE_DRIVE_LETTERS:
                        ErrInfo = "没有多余的盘符";
                        ErrCode2 = $"{ErrCode.ErrorCode.__MAINPROG__CREATEPE_NO_MORE_DRIVE_LETTERS}";
                        break;

                    case ErrCode.ErrorCode.__MAINPROG__CREATEPE__DISKPART__CREATE_PARTITION_FAIL:
                        ErrInfo = "无法通过DiskPart创建指定大小的分区";
                        ErrCode2 = $"{ErrCode.ErrorCode.__MAINPROG__CREATEPE_NO_MORE_DRIVE_LETTERS}";
                        break;

                    default:
                        ErrInfo = "未识别的错误";
                        break;

                }

                //报错返回
                MyError.ShowError(ErrCode.ErrorCode.__MAINPROG__CREATEPE_CREATE_PARTITION_FAIL, $"详细错误代码：{ErrCode2}\nGetLastError={Win32API.GetLastError()}\n{ErrInfo}");
                ReturnCode = ErrCode.ErrorCode.__MAINPROG__CREATEPE_CREATE_PARTITION_FAIL;
                return false;

            }
            LogProtocol.WriteLog.Info($"[CREATE_PE][DISKPART]   已创建分区\"{MainProgramPartitionLetter}\"");
            LogProtocol.WriteLog.Info($"[CREATE_PE][DISKPART]   被压缩分区\"{ShrinkPartitionLetter}\"");

            LogProtocol.WriteLog.Info("[CREATE_PE]   (1/2)正在创建并复制新的必要文件...");
            try
            {
                /* 创建标志文件 */
                File.Create($"{ReInstallPartitionLetter}:\\AZPEP_SysRi_Tool__ReInstallPartition").Close();
                File.Create($"{MainProgramPartitionLetter}:\\AZPEP_SysRi__MainProgramPartition").Close();
                File.Create($"{EfiPartitionLetter}:\\AZPEP_SysRi_Tool__EfiPartition").Close();
                if (ShrinkPartitionLetter != null)
                {
                    File.Create($"{ShrinkPartitionLetter}:\\AZPEP_SysRi_Tool__ShrinkPartition").Close();
                }

                /* 拷贝PE文件 */
                Directory.CreateDirectory($"{MainProgramPartitionLetter}:\\WinPE");
                File.Copy(wimFilePath, $"{MainProgramPartitionLetter}:\\WinPE\\boot.wim", true);
                File.Copy(sdiFilePath, $"{MainProgramPartitionLetter}:\\WinPE\\boot.sdi", true);


                /* 拷贝安装映像文件 */
                File.Copy(ImageFilePath, $"{MainProgramPartitionLetter}:\\Reinstall.{GetFileExName(ImageFilePath)}", true);
            }
            catch (Exception Err)
            {
                //报错返回

                MyError.ShowError(ErrCode.ErrorCode.__MAINPROG__CREATEPE_COPY_FILE_FAIL, $"详细错误信息：{Err.Message}", err: Err);
                PartitionParse((char)MainProgramPartitionLetter, ShrinkPartitionLetter);

                ReturnCode = ErrCode.ErrorCode.__MAINPROG__CREATEPE_COPY_FILE_FAIL;
                return false;
            }

            /* 创建PE启动项 */

            /* 首先，我们要删除之前创建的启动项，如果启动项之前创建过，程序将直接报错 */

            /*
            try
            {
                RunBCDedit($"/delete {GUID1} /cleanup");
                RunBCDedit($"/delete {GUID2} /cleanup");
            }
            finally
            {
                //我们不想让其报错，所以此处留空
            }
            */

            try
            {
                LogProtocol.WriteLog.Info("[CREATE_PE]   (2/2)正在创建PE启动项...");
                CreateStartupEntry(false, MainProgramPartitionLetter);
            }
            catch (Exception Err)
            {
                //报错返回
                MyError.ShowError(ErrCode.ErrorCode.__MAINPROG__CREATEPE_CREATE_PE_STARTUP_ENTRY_FAIL, $"详细错误信息：{Err.Message}", err: Err);
                PartitionParse((char)MainProgramPartitionLetter, ShrinkPartitionLetter);

                ReturnCode = ErrCode.ErrorCode.__MAINPROG__CREATEPE_CREATE_PE_STARTUP_ENTRY_FAIL;
                return false;
            }


            /* 创建配置文件 */
            FileStream FS = File.Create($"{MainProgramPartitionLetter}:\\config.ini");
            FS.Close();

            // 写入配置文件

            /*
             * 配置文件格式
             * [config]
             * InstallFile=(blahblah)
             * ImageIndex=(blahblah)
             * 
             * DisabledWindowsDefender=(0 or 1)
             * DisabledWindowsUpdate=(0 or 1)
             * SkipMicrosoftLogin=(0 or 1)
             * 注：
             * InstallFile                  映像文件路径
             * ImageIndex                   选中的映像文件索引
             * DisableWindowsDefender      禁用“Windows Defender”
             * DisableWindowsUpdate        禁用Windows自动更新
             * SkipMicrosoftLogin           跳过Microsoft登陆验证
             */

            Ini.Ini.IniFile IF = Ini.Ini.OpenIniFile($"{MainProgramPartitionLetter}:\\config.ini");
            IF.AddSection("config");
            IF.AddKeys("config", "InstallFile", $"Reinstall.{GetFileExName(ImageFilePath)}");
            IF.AddKeys("config", "ImageIndex", $"{ImageIndex}");
            IF.AddKeys("config", "DisableWindowsDefender", DisableWindowsDefender ? "1" : "0");
            IF.AddKeys("config", "DisableWindowsUpdate", DisableWindowsUpdate ? "1" : "0");
            IF.AddKeys("config", "SkipMicrosoftLogin", SkipMicrosoftLogin ? "1" : "0");

            ReturnCode = ErrCode.ErrorCode.SUCCESS;
            return true;

        }

        private static bool CreateStartupEntry(bool TryAgain, char? PartitionLetter)
        {
            try
            {
                /* 此段参考 “HotPE Client”中所给的代码 */

                RunBCDedit($"/create {GUID1} /d \"AZPEP SysRi Tool 重装工具\" /application osloader");
                RunBCDedit($"/create {GUID2} /d \"AZPEP SysRi Tool 重装工具\" /device");

                RunBCDedit($"/set {GUID2} ramdisksdidevice partition={PartitionLetter}:");
                RunBCDedit($"/set {GUID2} ramdisksdipath \\WinPE\\boot.sdi");

                RunBCDedit($"/set {GUID1} device ramdisk=[{PartitionLetter}:]\\WinPE\\boot.wim,{GUID2}");

                RunBCDedit($"/set {GUID1} path \\windows\\system32\\boot\\winload.efi");

                RunBCDedit($"/set {GUID1} description \"AZPEP SysRi Tool 重装工具\"");
                RunBCDedit($"/set {GUID1} osdevice ramdisk=[{PartitionLetter}:]\\WinPE\\boot.wim,{GUID2}");
                RunBCDedit($"/set {GUID1} systemroot \\windows");
                RunBCDedit($"/set {GUID1} detecthal Yes");
                RunBCDedit($"/set {GUID1} winpe Yes");

                RunBCDedit($"/displayorder {GUID1} /addlast");
            }
            catch (Exception Err)
            {
                try
                {
                    RunBCDedit($"/delete {GUID1} /cleanup");
                    RunBCDedit($"/delete {GUID2} /cleanup");
                }
                catch
                {
                    if (TryAgain)
                    {
                        MyError.ShowError(ErrCode.ErrorCode.__MAINPROG__CREATEPE_CREATE_PE_STARTUP_ENTRY_FAIL, $"详细错误信息：{Err.Message}", err: Err);
                        return false;
                    }
                    else CreateStartupEntry(true, PartitionLetter);
                }
            }

            return true;
        }
        /// <summary>
        /// 获取文件全名
        /// </summary>
        /// <param name="FullPath">文件完整路径</param>
        /// <returns></returns>
        public static string GetFileFullName(string FullPath)
        {
            string[] sL = FullPath.Split("\\");
            return sL[sL.Length - 1];
        }

        /// <summary>
        /// 获取文件扩展名
        /// </summary>
        /// <param name="FullPath">文件完整路径</param>
        /// <returns></returns>
        public static string GetFileExName(string FullPath)
        {
            string[] sL = FullPath.Split(".");
            return sL[sL.Length - 1];
        }
        private static bool RunBCDedit(
            string parameter
            )
        {
            Process p = new Process();
            //设置要启动的应用程序
            p.StartInfo.FileName = BcdEditFilePath;
            //是否使用操作系统shell启动
            p.StartInfo.UseShellExecute = false;
            // 接受来自调用程序的输入信息
            p.StartInfo.RedirectStandardInput = true;
            //输出信息
            p.StartInfo.RedirectStandardOutput = true;
            // 输出错误
            p.StartInfo.RedirectStandardError = true;
            //不显示程序窗口
            p.StartInfo.CreateNoWindow = true;

            p.StartInfo.Arguments = parameter;

            p.Start();

            StreamReader reader = p.StandardOutput;
            string output = reader.ReadToEnd();

            p.WaitForExit();
            Thread.Sleep(250);

            if (p.ExitCode != 0)
            {
                p.Close();

                throw new Exception(output);
                return false;
            }
            else
            {
                p.Close();
                return true;
            }
        }

        /// <summary>
        /// 在创建PE失败后处理分区
        /// </summary>
        /// <param name="PePartitionLetter"></param>
        /// <param name="ShrinkPartitionLetter"></param>
        public static void PartitionParse(
            char PePartitionLetter,
            char? ShrinkPartitionLetter
            )
        {
            DiskPart.DeleteVolume(PePartitionLetter);

            if (ShrinkPartitionLetter != null)
            {
                DiskPart.ExtendVolume((char)ShrinkPartitionLetter);
            }
        }
    }
}
